package net.cyue.ort.llm.sampling;

import net.cyue.ort.llm.data.Token;

import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

/**
 * 随机采样策略
 * 支持温度、Top-K、Top-P、重复惩罚等常见算法
 */
public class RandomSamplingStrategy implements SamplingStrategy {
    private final Random random;

    public RandomSamplingStrategy() {
        this(new SecureRandom());
    }

    public RandomSamplingStrategy(Random random) {
        this.random = random;
    }

    @Override
    public Token selectNextToken(float[] logits, List<Long> context, SamplingConfig config) {
        List<ProbabilityToken> distribution = buildDistribution(logits, context, config);
        if (distribution.isEmpty()) {
            return new Token(0L);
        }
        double target = random.nextDouble();
        double cumulative = 0d;
        for (ProbabilityToken token : distribution) {
            cumulative += token.probability();
            if (target <= cumulative) {
                return token.token();
            }
        }
        return distribution.get(distribution.size() - 1).token();
    }

    @Override
    public List<Token> selectTopNTokens(float[] logits, List<Long> context, SamplingConfig config, int topN) {
        List<ProbabilityToken> distribution = buildDistribution(logits, context, config);
        int limit = Math.min(topN, distribution.size());
        List<Token> result = new ArrayList<>(limit);
        for (int i = 0; i < limit; i++) {
            result.add(distribution.get(i).token());
        }
        return result;
    }

    private List<ProbabilityToken> buildDistribution(float[] logits, List<Long> context, SamplingConfig config) {
        if (logits == null || logits.length == 0) {
            return List.of();
        }
        SamplingConfig effectiveConfig = config != null ? config : SamplingConfig.defaultConfig();
        double[] adjusted = SamplingUtils.toDoubleArray(logits);
        SamplingUtils.applyRepetitionPenalty(adjusted, context, effectiveConfig.getRepetitionPenalty());
        float temperature = effectiveConfig.getTemperature();
        SamplingUtils.applyTemperature(adjusted, temperature <= 0f ? 1f : temperature);

        List<SamplingUtils.Candidate> candidates = SamplingUtils.buildCandidates(
            adjusted,
            effectiveConfig.getTopK(),
            effectiveConfig.getTopP()
        );

        List<ProbabilityToken> distribution = new ArrayList<>(candidates.size());
        for (SamplingUtils.Candidate candidate : candidates) {
            Token token = new Token(candidate.tokenId(), candidate.logit());
            distribution.add(new ProbabilityToken(token, candidate.probability()));
        }
        return distribution;
    }

    private record ProbabilityToken(Token token, double probability) {
    }
}


